#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
新建数据集对话框
"""

import os
from PyQt5.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QFormLayout,
    QLabel, QLineEdit, QComboBox, QPushButton,
    QFileDialog, QMessageBox, QGroupBox
)
from PyQt5.QtCore import Qt, QSize

class NewDatasetDialog(QDialog):
    """新建数据集对话框"""
    
    def __init__(self, config, parent=None):
        """
        初始化对话框
        
        Args:
            config (dict): 应用配置
            parent (QWidget, optional): 父窗口
        """
        super().__init__(parent)
        
        self.config = config
        
        # 设置对话框属性
        self.setWindowTitle("新建数据集")
        self.setMinimumSize(500, 300)
        
        # 初始化UI
        self._setup_ui()
        
        # 初始化输入
        self._initialize_input()
    
    def _setup_ui(self):
        """设置界面"""
        # 主布局
        main_layout = QVBoxLayout(self)
        
        # 表单布局
        form_layout = QFormLayout()
        
        # 数据集名称
        self.name_edit = QLineEdit()
        self.name_edit.setPlaceholderText("输入数据集名称")
        form_layout.addRow("数据集名称:", self.name_edit)
        
        # 数据集保存路径
        path_layout = QHBoxLayout()
        self.path_edit = QLineEdit()
        self.path_edit.setReadOnly(True)
        self.path_edit.setPlaceholderText("选择数据集保存路径")
        self.browse_button = QPushButton("浏览...")
        self.browse_button.clicked.connect(self._browse_path)
        path_layout.addWidget(self.path_edit)
        path_layout.addWidget(self.browse_button)
        form_layout.addRow("保存路径:", path_layout)
        
        # 算法类型
        self.algorithm_combo = QComboBox()
        self.algorithm_combo.addItem("图像分类", "classification")
        self.algorithm_combo.addItem("目标检测", "detection")
        form_layout.addRow("算法类型:", self.algorithm_combo)
        
        # 添加表单到主布局
        form_group = QGroupBox("基本信息")
        form_group.setLayout(form_layout)
        main_layout.addWidget(form_group)
        
        # 算法详细信息组
        self.algorithm_group = QGroupBox("算法参数")
        self.algorithm_layout = QFormLayout(self.algorithm_group)
        
        # 分类模型选择
        self.model_combo = QComboBox()
        for model in self.config['algorithms']['classification']['available']:
            self.model_combo.addItem(model)
        self.algorithm_layout.addRow("模型:", self.model_combo)
        
        main_layout.addWidget(self.algorithm_group)
        
        # 添加弹性空间
        main_layout.addStretch(1)
        
        # 底部按钮布局
        button_layout = QHBoxLayout()
        self.cancel_button = QPushButton("取消")
        self.cancel_button.clicked.connect(self.reject)
        self.create_button = QPushButton("创建")
        self.create_button.setDefault(True)
        self.create_button.clicked.connect(self._handle_create)
        button_layout.addWidget(self.cancel_button)
        button_layout.addWidget(self.create_button)
        
        main_layout.addLayout(button_layout)
    
    def _initialize_input(self):
        """初始化输入"""
        # 设置默认路径
        default_path = self.config['paths']['datasets']
        if os.path.exists(default_path):
            self.path_edit.setText(default_path)
        
        # 设置默认算法类型
        default_algo = self.config['algorithms']['classification']['default']
        index = self.model_combo.findText(default_algo)
        if index >= 0:
            self.model_combo.setCurrentIndex(index)
        
        # 连接算法类型变更信号
        self.algorithm_combo.currentIndexChanged.connect(self._on_algorithm_changed)
    
    def _browse_path(self):
        """浏览保存路径"""
        current_path = self.path_edit.text()
        if not current_path:
            current_path = os.path.expanduser("~")
        
        dir_path = QFileDialog.getExistingDirectory(
            self, 
            "选择数据集保存路径", 
            current_path
        )
        
        if dir_path:
            self.path_edit.setText(dir_path)
    
    def _on_algorithm_changed(self, index):
        """
        算法类型变更处理
        
        Args:
            index (int): 当前索引
        """
        # 获取当前算法类型
        algo_type = self.algorithm_combo.currentData()
        
        # 清除当前算法参数布局
        while self.algorithm_layout.rowCount() > 0:
            self.algorithm_layout.removeRow(0)
        
        # 根据算法类型设置参数
        if algo_type == "classification":
            # 分类模型选择
            self.model_combo = QComboBox()
            for model in self.config['algorithms']['classification']['available']:
                self.model_combo.addItem(model)
            
            # 设置默认模型
            default_model = self.config['algorithms']['classification']['default']
            index = self.model_combo.findText(default_model)
            if index >= 0:
                self.model_combo.setCurrentIndex(index)
                
            self.algorithm_layout.addRow("模型:", self.model_combo)
        
        elif algo_type == "detection":
            # 检测模型选择
            self.model_combo = QComboBox()
            for model in self.config['algorithms']['detection']['available']:
                self.model_combo.addItem(model)
            
            # 设置默认模型
            default_model = self.config['algorithms']['detection']['default']
            index = self.model_combo.findText(default_model)
            if index >= 0:
                self.model_combo.setCurrentIndex(index)
                
            self.algorithm_layout.addRow("模型:", self.model_combo)
    
    def _handle_create(self):
        """处理创建数据集"""
        # 验证输入
        name = self.name_edit.text().strip()
        path = self.path_edit.text().strip()
        algo_type = self.algorithm_combo.currentData()
        model = self.model_combo.currentText()
        
        # 验证数据集名称
        if not name:
            QMessageBox.warning(self, "验证错误", "请输入数据集名称")
            self.name_edit.setFocus()
            return
        
        # 验证保存路径
        if not path:
            QMessageBox.warning(self, "验证错误", "请选择数据集保存路径")
            return
        
        if not os.path.exists(path):
            try:
                os.makedirs(path, exist_ok=True)
            except Exception as e:
                QMessageBox.critical(self, "路径错误", f"无法创建路径: {e}")
                return
        
        # 设置结果
        self.result = {
            "name": name,
            "path": path,
            "algorithm_type": algo_type,
            "model": model
        }
        
        # 接受对话框
        self.accept()
    
    def get_result(self):
        """
        获取对话框结果
        
        Returns:
            dict: 包含数据集配置的字典
        """
        if hasattr(self, 'result'):
            return self.result
        return None 